#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Example usage of the modified Hi-C dataset with NPZ format and windowed sampling
"""

import numpy as np
import torch
from Squidiff.hic_datasets import create_hic_dataloader, HiCDataset

def create_sample_data():
    """
    Create sample Hi-C data in NPZ format for demonstration
    """
    # Create sample Hi-C matrices for different chromosomes
    chr_sizes = {'chr1': 1000, 'chr2': 800, 'chr3': 600}
    
    # Create two timepoints
    timepoints = ['G1_top400', 'late_S1_top400'] 
    
    for timepoint in timepoints:
        # Create directory
        import os
        os.makedirs(f'sample_data/{timepoint}', exist_ok=True)
        
        # Create 3 sample cells per timepoint
        for cell_id in range(3):
            filename = f'sample_data/{timepoint}/cell_{cell_id:03d}.npz'
            
            # Create Hi-C matrices for each chromosome
            matrices = {}
            for chr_name, size in chr_sizes.items():
                # Create a synthetic Hi-C matrix with distance decay
                matrix = np.zeros((size, size))
                for i in range(size):
                    for j in range(size):
                        distance = abs(i - j)
                        if distance == 0:
                            contact = np.random.poisson(100)
                        elif distance < 50:
                            contact = np.random.poisson(max(1, 50 - distance))
                        else:
                            contact = np.random.poisson(1) if np.random.random() < 0.1 else 0
                        matrix[i, j] = contact
                        matrix[j, i] = contact  # Make symmetric
                
                # Add some noise to differentiate timepoints
                if timepoint == 'late_S1_top400':
                    matrix = matrix * (1 + 0.2 * np.random.random(matrix.shape))
                
                matrices[chr_name] = matrix.astype(np.float32)
            
            # Save as NPZ file
            np.savez(filename, **matrices)
            print(f"Created {filename}")

def test_dataset():
    """
    Test the modified dataset with different configurations
    """
    print("Testing Hi-C Dataset with NPZ format...")
    
    # Test configuration 1: Basic usage
    print("\n1. Basic NPZ dataset loading:")
    dataset = HiCDataset(
        data_dir='sample_data',
        timepoints=['G1_top400', 'late_S1_top400'],
        file_format='npz',
        chromosome='chr1',
        window_size=128,  # Custom window size
        random_window=True
    )
    
    print(f"Dataset length: {len(dataset)}")
    
    # Get a sample
    sample = dataset[0]
    print(f"Sample keys: {sample.keys()}")
    print(f"Hi-C matrix shape: {sample['hic'].shape}")
    print(f"Time value: {sample['time'].item()}")
    print(f"Cell ID: {sample['cell_id']}")
    print(f"Timepoint: {sample['timepoint']}")
    
    # Test configuration 2: Different window sizes
    print("\n2. Testing different window sizes:")
    for window_size in [64, 128, 256, 512]:
        try:
            dataset = HiCDataset(
                data_dir='sample_data',
                timepoints=['G1_top400', 'late_S1_top400'],
                file_format='npz',
                chromosome='chr1',
                window_size=window_size,
                random_window=True
            )
            sample = dataset[0]
            print(f"Window size {window_size}: Success, matrix shape {sample['hic'].shape}")
        except Exception as e:
            print(f"Window size {window_size}: Failed - {e}")
    
    # Test configuration 3: Different chromosomes
    print("\n3. Testing different chromosomes:")
    for chr_name in ['chr1', 'chr2', 'chr3']:
        try:
            dataset = HiCDataset(
                data_dir='sample_data',
                timepoints=['G1_top400', 'late_S1_top400'],
                file_format='npz',
                chromosome=chr_name,
                window_size=128,
                random_window=True
            )
            sample = dataset[0]
            print(f"Chromosome {chr_name}: Success, matrix shape {sample['hic'].shape}")
        except Exception as e:
            print(f"Chromosome {chr_name}: Failed - {e}")

def test_dataloader():
    """
    Test the dataloader functionality
    """
    print("\n4. Testing DataLoader:")
    
    dataloader = create_hic_dataloader(
        data_dir='sample_data',
        batch_size=2,
        timepoints=['G1_top400', 'late_S1_top400'],
        file_format='npz',
        chromosome='chr1',
        window_size=128,
        random_window=True,
        num_workers=0  # Use 0 for debugging
    )
    
    print(f"DataLoader created successfully")
    print(f"Number of batches: {len(dataloader)}")
    
    # Test one batch
    for batch_idx, batch in enumerate(dataloader):
        print(f"Batch {batch_idx}:")
        print(f"  Hi-C shape: {batch['hic'].shape}")
        print(f"  Time shape: {batch['time'].shape}")
        print(f"  Cell IDs: {batch['cell_id']}")
        print(f"  Timepoints: {batch['timepoint']}")
        if batch_idx >= 1:  # Only show first 2 batches
            break

def main():
    """Main function"""
    print("Hi-C Dataset NPZ Format Example")
    print("=" * 50)
    
    # Step 1: Create sample data
    print("Creating sample data...")
    create_sample_data()
    
    # Step 2: Test dataset
    test_dataset()
    
    # Step 3: Test dataloader
    test_dataloader()
    
    print("\nExample completed successfully!")
    print("\nUsage Summary:")
    print("- Data format: NPZ files containing chromosome-specific Hi-C matrices")
    print("- Window sampling: Extract square windows from full chromosome matrices")
    print("- Chromosome-specific: Only intra-chromosome interactions are considered")
    print("- Customizable: Window size, normalization, augmentation options available")

if __name__ == "__main__":
    main() 